#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-08-07

"""
【节点名称】：
    CameraNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    无
【发布类型】：
    sensor_msgs::CompressedImage （输出图像）
    sensor_msgs::CameraCalibration （输出标定信息）
【构造参数说明】：
    parameter_file (str): 全局参数文件
【节点参数】：
    resize (list): 图像重采样分辨率
    camera_id (int): 相机ID号
【备注】：
    无
"""

import threading
import time
import cv2
import argparse
from typing import Union
from spirems import Publisher, cvimg2sms, def_msg
from spirecv.base.BaseNode import BaseNode


class CameraNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        specified_output_topic: str = None,
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file
        )
        
        output_url = '/' + job_name + '/sensor/calibration_info'
        if specified_output_topic is not None:
            output_url = specified_output_topic
            
        self._image_writer = Publisher(
            '/' + job_name + '/sensor/image_raw', 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )
        
        self._calibration_writer = Publisher(
            output_url, 'sensor_msgs::CameraCalibration',
            ip=ip, port=port
        )

        # camera configure
        self.use_v4l2 = self.get_param("use_v4l2", 0)
        self.camera_id = self.get_param("camera_id", 0)
        self.width = self.get_param("width", 640)
        self.height = self.get_param("height", 480)
        self.frame_id = self.get_param("frame_id", "camera")
        self.fourcc = self.get_param("fourcc", "MJPG")
        self.fps = self.get_param("fps", 30)
        
        # camera calibration
        
        self.camera_matrix =  self.get_param("camera_matrix", [
            7.9379415710551370e+02, 0, 2.9783879354295328e+02, 0,
            7.9491985564466654e+02, 3.0942416136837386e+02, 0, 0, 1])
        self.distortion_coefficients = self.get_param("distortion_coefficients",[
            2.0950200339181715e-01, -1.1587468096518483e+00,
            5.5342063671841328e-03, 2.2214393775334758e-04, 1.7127431916651392e+00])
        self.projection = self.get_param("projection", [
            7.9379415710551370e+02, 0, 2.9783879354295328e+02, 0, 0,
            7.9491985564466654e+02, 3.0942416136837386e+02, 0, 0, 0, 1, 0])
        if self.use_v4l2 > 0:
            self.capture = cv2.VideoCapture(self.camera_id, cv2.CAP_V4L2)
        else:
            self.capture = cv2.VideoCapture(self.camera_id)
        self.capture.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
        self.capture.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
        self.capture.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*self.fourcc))
        if self.fps > 0:
            self.capture.set(cv2.CAP_PROP_FPS, self.fps)
        self.img_id = 0
        self.last_check_time = time.time() - 1     
        self.interval = 1
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_writer.kill()
        self._image_writer.join()
        self._calibration_writer.kill()
        self._calibration_writer.join()

    def run(self):
        while self.is_running():
            ret, frame = self.capture.read()
            if not ret:
                continue
            frame = cv2.resize(frame, (self.width, self.height))
            msg = cvimg2sms(frame)
            self.img_id += 1
            msg['img_id'] = self.img_id
            self._image_writer.publish(msg)
            # cv2.imshow('img', frame)
            # cv2.waitKey(5)
            
            current_time = time.time()
            if current_time - self.last_check_time >= self.interval:
                msg_calb = def_msg('sensor_msgs::CameraCalibration')
                msg_calb['frame_id'] = self.frame_id
                msg_calb['width'] = self.width
                msg_calb['height'] = self.height
                msg_calb['K'] = self.camera_matrix
                msg_calb['D'] = self.distortion_coefficients
                msg_calb['P'] = self.projection
                # publish calibration info in 1 Hz
                self._calibration_writer.publish(msg_calb)
                # current time
                self.last_check_time = current_time
            
            
            self._image_writer.publish(msg)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        'config',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    args = parser.parse_args()
    print("config:", args.config)
    print("job-name:", args.job_name)
    cam = CameraNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port)
    cam.join()

